"""
WebSocket Manager

Manages WebSocket connections for real-time communication in the web interface.
"""

import asyncio
import json
import logging
from typing import List, Dict, Any, Optional

try:
    from fastapi import WebSocket, WebSocketDisconnect
    FASTAPI_AVAILABLE = True
except ImportError:
    FASTAPI_AVAILABLE = False
    WebSocket = None

logger = logging.getLogger(__name__)


class WebSocketManager:
    """
    Manages WebSocket connections for real-time communication.
    
    Handles multiple concurrent connections and message routing between
    clients and the AI Historical Simulation Platform.
    """
    
    def __init__(self, platform):
        """
        Initialize WebSocket manager.
        
        Args:
            platform: AI Historical Simulation Platform instance
        """
        self.platform = platform
        self.active_connections: List[WebSocket] = []
        self.connection_sessions: Dict[WebSocket, Dict[str, Any]] = {}
        
        logger.info("WebSocket manager initialized")
    
    async def connect(self, websocket: WebSocket) -> None:
        """
        Accept a new WebSocket connection.
        
        Args:
            websocket: WebSocket connection to accept
        """
        if not FASTAPI_AVAILABLE:
            logger.error("FastAPI not available for WebSocket connections")
            return
        
        await websocket.accept()
        self.active_connections.append(websocket)
        
        # Initialize session data for this connection
        self.connection_sessions[websocket] = {
            'user_id': f'ws_user_{len(self.active_connections)}',
            'current_figure': None,
            'session_id': f'ws_session_{self.platform.get_timestamp()}_{len(self.active_connections)}',
            'connected_at': self.platform.get_timestamp()
        }
        
        logger.info(f"WebSocket connection accepted. Total connections: {len(self.active_connections)}")
        
        # Send welcome message
        await self.send_message(websocket, {
            'type': 'welcome',
            'message': 'Connected to AI Historical Simulation Platform',
            'session_id': self.connection_sessions[websocket]['session_id']
        })
    
    def disconnect(self, websocket: WebSocket) -> None:
        """
        Remove a WebSocket connection.
        
        Args:
            websocket: WebSocket connection to remove
        """
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        
        if websocket in self.connection_sessions:
            del self.connection_sessions[websocket]
        
        logger.info(f"WebSocket connection removed. Total connections: {len(self.active_connections)}")
    
    async def send_message(self, websocket: WebSocket, message: Dict[str, Any]) -> None:
        """
        Send a message to a specific WebSocket connection.
        
        Args:
            websocket: Target WebSocket connection
            message: Message to send
        """
        try:
            await websocket.send_text(json.dumps(message))
        except Exception as e:
            logger.error(f"Error sending WebSocket message: {e}")
            # Remove disconnected websocket
            self.disconnect(websocket)
    
    async def broadcast_message(self, message: Dict[str, Any], exclude: Optional[WebSocket] = None) -> None:
        """
        Broadcast a message to all connected WebSocket clients.
        
        Args:
            message: Message to broadcast
            exclude: Optional WebSocket to exclude from broadcast
        """
        disconnected = []
        
        for websocket in self.active_connections:
            if websocket != exclude:
                try:
                    await websocket.send_text(json.dumps(message))
                except Exception as e:
                    logger.error(f"Error broadcasting to WebSocket: {e}")
                    disconnected.append(websocket)
        
        # Clean up disconnected websockets
        for ws in disconnected:
            self.disconnect(ws)
    
    async def handle_message(self, websocket: WebSocket, data: Dict[str, Any]) -> None:
        """
        Handle incoming WebSocket message.
        
        Args:
            websocket: Source WebSocket connection
            data: Message data
        """
        try:
            message_type = data.get('type')
            session_data = self.connection_sessions.get(websocket, {})
            
            if message_type == 'get_figures':
                await self._handle_get_figures(websocket)
                
            elif message_type == 'select_figure':
                await self._handle_select_figure(websocket, data, session_data)
                
            elif message_type == 'chat_message':
                await self._handle_chat_message(websocket, data, session_data)
                
            elif message_type == 'get_metrics':
                await self._handle_get_metrics(websocket)
                
            elif message_type == 'ping':
                await self.send_message(websocket, {'type': 'pong'})
                
            else:
                await self.send_message(websocket, {
                    'type': 'error',
                    'message': f'Unknown message type: {message_type}'
                })
                
        except Exception as e:
            logger.error(f"Error handling WebSocket message: {e}")
            await self.send_message(websocket, {
                'type': 'error',
                'message': 'Internal server error'
            })
    
    async def _handle_get_figures(self, websocket: WebSocket) -> None:
        """Handle request for available figures."""
        try:
            figures = self.platform.figure_manager.list_available_figures()
            await self.send_message(websocket, {
                'type': 'figures_list',
                'figures': figures
            })
        except Exception as e:
            logger.error(f"Error getting figures: {e}")
            await self.send_message(websocket, {
                'type': 'error',
                'message': 'Failed to load figures'
            })
    
    async def _handle_select_figure(self, websocket: WebSocket, data: Dict[str, Any], 
                                   session_data: Dict[str, Any]) -> None:
        """Handle figure selection."""
        try:
            figure_name = data.get('figure')
            if not figure_name:
                await self.send_message(websocket, {
                    'type': 'error',
                    'message': 'Figure name is required'
                })
                return
            
            # Load the figure
            figure = await self.platform.load_historical_figure(figure_name)
            
            # Update session data
            session_data['current_figure'] = figure_name
            session_data['figure_personality_id'] = figure.personality_id
            
            await self.send_message(websocket, {
                'type': 'figure_selected',
                'figure': figure_name,
                'message': f'You are now chatting with {figure_name}'
            })
            
        except Exception as e:
            logger.error(f"Error selecting figure: {e}")
            await self.send_message(websocket, {
                'type': 'error',
                'message': f'Failed to select figure: {str(e)}'
            })
    
    async def _handle_chat_message(self, websocket: WebSocket, data: Dict[str, Any], 
                                  session_data: Dict[str, Any]) -> None:
        """Handle chat message."""
        try:
            message = data.get('message')
            figure_name = data.get('figure') or session_data.get('current_figure')
            
            if not message:
                await self.send_message(websocket, {
                    'type': 'error',
                    'message': 'Message is required'
                })
                return
            
            if not figure_name:
                await self.send_message(websocket, {
                    'type': 'error',
                    'message': 'No figure selected'
                })
                return
            
            # Load figure if needed
            figure = await self.platform.load_historical_figure(figure_name)
            
            # Generate response
            response = await self.platform.generate_response(
                figure.personality_id,
                message,
                session_data['session_id']
            )
            
            await self.send_message(websocket, {
                'type': 'chat_response',
                'figure': figure_name,
                'user_message': message,
                'response': response,
                'timestamp': self.platform.get_timestamp()
            })
            
        except Exception as e:
            logger.error(f"Error processing chat message: {e}")
            await self.send_message(websocket, {
                'type': 'error',
                'message': 'Failed to process message'
            })
    
    async def _handle_get_metrics(self, websocket: WebSocket) -> None:
        """Handle request for system metrics."""
        try:
            metrics = self.platform.get_performance_metrics()
            health = self.platform.get_system_health()
            
            await self.send_message(websocket, {
                'type': 'metrics_update',
                'metrics': metrics,
                'health': health
            })
            
        except Exception as e:
            logger.error(f"Error getting metrics: {e}")
            await self.send_message(websocket, {
                'type': 'error',
                'message': 'Failed to load metrics'
            })
    
    def get_connection_count(self) -> int:
        """Get the number of active connections."""
        return len(self.active_connections)
    
    def get_connection_info(self) -> List[Dict[str, Any]]:
        """Get information about all active connections."""
        info = []
        for i, websocket in enumerate(self.active_connections):
            session_data = self.connection_sessions.get(websocket, {})
            info.append({
                'connection_id': i,
                'user_id': session_data.get('user_id', 'unknown'),
                'session_id': session_data.get('session_id', 'unknown'),
                'current_figure': session_data.get('current_figure', 'none'),
                'connected_at': session_data.get('connected_at', 'unknown')
            })
        return info